from paddle_serving_client import Client
from paddle_serving_app.reader import *
from PIL import Image, ImageDraw
import os
import numpy as np
import time
import cv2
import sys
py_version = sys.version_info[0]
if py_version == 2:
    import urllib
else:
    import urllib.request as urllib

# 将服务器返回数据处理成ｊｓｏｎ
class Postprocess(object):
    def __init__(self, image_name,label_file, output_dir, resize_shape=None, threshold=0.5, draw_bbox_image=False):
        self.image_name = image_name
        self.output_dir = output_dir
        self.label_file = label_file
        self.label_list = []
        self.threshold = threshold
        self.resize_shape = resize_shape
        self.draw_bbox_image = draw_bbox_image
        self.label_list = label_file
        self.clsid2catid = {i: i for i in range(len(self.label_list))}
        self.catid2name = {i: name for i, name in enumerate(self.label_list)}

    def _offset_to_lengths(self, lod):
        offset = lod[0]
        lengths = [offset[i + 1] - offset[i] for i in range(len(offset) - 1)]
        return [lengths]

    def _bbox2out(self, results, clsid2catid, is_bbox_normalized=True, ):
        xywh_res = []
        for t in results:
            bboxes = t['bbox'][0]
            lengths = t['bbox'][1][0]
            if bboxes.shape == (1, 1) or bboxes is None:
                continue

            k = 0
            for i in range(len(lengths)):
                num = lengths[i]
                for j in range(num):
                    dt = bboxes[k]
                    clsid, score, xmin, ymin, xmax, ymax = dt.tolist()
                    catid = (clsid2catid[int(clsid)])
                    if self.threshold > score:
                        continue
                    if is_bbox_normalized:
                        # xmin, ymin, xmax, ymax = \
                        #     self.clip_bbox([xmin, ymin, xmax, ymax])
                        w = xmax - xmin
                        h = ymax - ymin
                        im_shape = [608,608]
                        im_height, im_width = int(im_shape[0]), int(im_shape[1])
                        xmin /= im_width
                        ymin /= im_height
                        w /= im_width
                        h /= im_height
                    else:
                        w = xmax - xmin + 1
                        h = ymax - ymin + 1

                    bbox = [xmin, ymin, w, h]
                    coco_res = {
                        'name':self.image_name,
                        'type': catid,
                        'rect': bbox,
                        #'score': score
                    }
                    xywh_res.append(coco_res)
                    k += 1
        return xywh_res

    def _get_bbox_result(self, fetch_map, fetch_name, clsid2catid):
        result = {}
        is_bbox_normalized = True
        output = fetch_map[fetch_name]
        lod = [fetch_map[fetch_name + '.lod']]
        lengths = self._offset_to_lengths(lod)
        np_data = np.array(output)
        result['bbox'] = (np_data, lengths)
        result['im_id'] = np.array([[0]])

        bbox_results = self._bbox2out([result], clsid2catid, is_bbox_normalized)
        return bbox_results

    def color_map(self, num_classes):
        color_map = num_classes * [0, 0, 0]
        for i in range(0, num_classes):
            j = 0
            lab = i
            while lab:
                color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
                color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
                color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
                j += 1
                lab >>= 3
        color_map = np.array(color_map).reshape(-1, 3)
        return color_map

    def draw_bbox(self, image, catid2name, bboxes, threshold, color_list):
        """
        draw bbox on image
        """
        draw = ImageDraw.Draw(image)

        for dt in np.array(bboxes):
            catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
            if score < threshold:
                continue

            xmin, ymin, w, h = bbox
            xmax = xmin + w
            ymax = ymin + h

            img_w, img_h = image.size
            if self.resize_shape is not None:
                xmin = xmin * img_w / self.resize_shape[0]
                xmax = xmax * img_w / self.resize_shape[0]
                ymin = ymin * img_h / self.resize_shape[1]
                ymax = ymax * img_h / self.resize_shape[1]

            color = tuple(color_list[catid])

            # draw bbox
            draw.line(
                [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
                 (xmin, ymin)],
                width=2,
                fill=color)

            # draw label
            text = "{} {:.2f}".format(catid2name[catid], score)
            tw, th = draw.textsize(text)
            draw.rectangle(
                [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
            draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))

        return image

    def visualize(self, infer_img, bbox_results, catid2name, num_classes):
        image = Image.open(infer_img).convert('RGB')
        color_list = self.color_map(num_classes)
        image = self.draw_bbox(image, self.catid2name, bbox_results, self.threshold,
                               color_list)
        image_path = os.path.split(infer_img)[-1]
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
        out_path = os.path.join(self.output_dir, image_path)
        image.save(out_path, quality=95)

    # draw_bbox_image:将有标记框的图像画出来，默认不画图，只拿ｊｓｏｎ
    def __call__(self, image_with_bbox):
        fetch_name = ""
        for key in image_with_bbox:
            if key == "image":
                continue
            if ".lod" in key:
                continue
            fetch_name = key
        bbox_result = self._get_bbox_result(image_with_bbox, fetch_name,
                                            self.clsid2catid)
        if (self.draw_bbox_image):
            if os.path.isdir(self.output_dir) is False:
                os.mkdir(self.output_dir)
            self.visualize(image_with_bbox["image"], bbox_result, self.catid2name,
                           len(self.label_list))
        # bbox_file = os.path.join(self.output_dir, 'bbox.json')
        # with open(bbox_file, 'w') as f:
        #     json.dump(bbox_result, f, indent=4)
        return bbox_result

def url2img(img_url):
    resp = urllib.urlopen(img_url)
    sample = resp.read()
    data = np.fromstring(sample, np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return img
def path2img(img_path):
    if py_version == 2:
        fin = open(img_path)
    else:
        fin = open(img_path, "rb")
    sample = fin.read()
    data = np.fromstring(sample, np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
    '''
    img = cv2.imread(img_path, -1)
    channels = img.shape[2]
    ori_h = img.shape[0]
    ori_w = img.shape[1]
    '''
    return img

def paddle_client_create(model):
    '''
    model = suit :工作服分类模型;
    model = person_pig_car :人车猪检测模型;
    return ： client, pre_seq
    '''
    if model == 'suit':
        model = 1
    if model == 'person_pig_car':
        model = 2
    model_dict = {1: ["192.168.1.234:9393", 'serving/ppcls_client_conf/serving_client_conf.prototxt'],
                  2: ["192.168.1.234:9292", "/home/bcw/paddle/serving_client/serving_client_conf.prototxt"]}
    config = model_dict[model]
    client = Client()
    client.load_client_config(config[1])
    client.connect([config[0]])

    suit_prework = Sequential([Resize(256), CenterCrop(224), RGB2BGR(), Transpose((2, 0, 1)), \
                                Div(255), Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], True)])
    person_pig_car_prework = Sequential(
        [Resize((608, 608)), BGR2RGB(), Transpose((2, 0, 1)), \
         Div(255.0), Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], True)])
    if model == 1 :
        prework = suit_prework

    if model == 2 :
        prework = person_pig_car_prework

    return prework, client


def paddle_model_suit(image,image_type, prework, client, threshold):
    '''
    image:图片路径或者url；
    image_type: 0,url,1,图片,2，缓存；
    prework:预处理序列；
    client:client;
    threshold:阈值，eg：0.5;
    return：服务器返回结果 0:没穿 1:红 2:橙 3:黄 4:绿
    '''
    label_dict = (0, 1, 2, 3, 4)
    image_file = image
    if image_type == 1:
        image_file = path2img(image_file)
    if image_type == 0:
        image_file = url2img(image_file)
    img = prework(image_file)
    fetch_map = client.predict(feed={"image": img}, fetch=["prediction"])
    prob = max(fetch_map["prediction"][0])
    label = label_dict[fetch_map["prediction"][0].tolist().index(prob)]
    if prob > threshold:
        return label


def paddle_model_person_pig_car(image, image_type,prework, client, threshold,out=False):
    '''
    image:图片路径或者url；
    image_type: 0,url,1,图片；
    prework:预处理序列；
    client:client;
    threshold:阈值，eg：0.5;
    return：服务器返回结果 0:车，1：人，2：猪
    '''
    label_dict = (1, 2, 3)
    image_file = image
    if image_type:
        image_file = path2img(image_file)
    else:
        image_file = url2img(image_file)
    img = prework(image_file)
    fetch_map = client.predict(
        feed={
            "image": img,
            "im_size": np.array(list(img.shape[1:])),
        },
        fetch=["matrix_nms_0.tmp_0"])
    post = Postprocess(1,label_dict, "output", [608, 608], threshold)
    if out :
        post2 = RCNNPostprocess("label_list.txt","output",[608,608])
        fetch_map["image"] = image
        post2(fetch_map)
    result = post(fetch_map)
    return result,image_file

def paddle_get_suit(item,img):
    rect = item['rect'];
    img_w = img.shape[1]
    img_h = img.shape[0]
    y0 = int(rect[1]*img_h)
    y1 = int((rect[1]+rect[3])*img_h)
    x0 = int(rect[0]*img_w)
    x1 = int((rect[0]+rect[2])*img_w)
    img1 = img[y0:y1,x0:x1]
    cv2.imshow('1',img1)
    cv2.waitKey(0)
    return img1
